import os
import gradio as gr
import agentscope
from agentscope.agents import UserAgent
from agentscope.service import ServiceToolkit, execute_python_code, execute_shell_command
from agentscope.agents.react_agent import ReActAgent
from agentscope.message.msg import Msg

# 初始化模型配置
model_configs = [
    {
        "model_type": "openai_chat",
        "config_name": "qwen2.5-7b-local",
        "model_name": "qwen2.5:7b-instruct-q4_0",
        "client_args": {
            "base_url": "http://127.0.0.1:11434/v1",
        },
        "api_key": "Empty",
        "generate_args": {"temperature": 0.2},
        "stream": True,
    },
    {
        "model_type": "openai_chat",
        "config_name": "qwen2.5-7b",
        "model_name": "Qwen/Qwen2.5-7B-Instruct",
        "client_args": {
            "base_url": "https://api-inference.modelscope.cn/v1",
        },
        "api_key": os.environ.get("MODELSCOPE_ACCESS_TOKEN"),
        "generate_args": {"temperature": 0.2},
        "stream": True,
    },
    {
        "model_type": "openai_chat",
        "config_name": "deepseek-r1:7b",
        "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
        "client_args": {
            "base_url": "https://api.siliconflow.cn/v1",
        },
        "api_key": os.environ.get("SILICON_API_KEY"),
        "generate_args": {"temperature": 0.6},
        "stream": True,
    }
]

# agentscope.init(model_configs=model_configs, studio_url="http://127.0.0.1:5000/")
agentscope.init(model_configs=model_configs)


# 准备工具
service_toolkit = ServiceToolkit()
service_toolkit.add(execute_python_code)
service_toolkit.add(execute_shell_command)

# 创建代理
agent = ReActAgent(
    name="assistant",
    # model_config_name="qwen2.5-7b-local",
    model_config_name="qwen2.5-7b",
    # model_config_name="deepseek-r1:7b",
    verbose=True,
    service_toolkit=service_toolkit,
    max_iters=3,
    sys_prompt="你是一个擅长使用命令行工具以及代码生成的电脑助手，你操作的电脑是Windows11,你可以通过命令查询，如wmic命令查询系统信息来回答用户问题，以及通过python代码生成和调用代码执行工具来回答复杂问题的能力！"
)
user = UserAgent(name="User", input_hint="User Input ('exit' to quit): ")

# 定义对话函数
def chat(message, history):
    history = history or []
    msg = Msg(name="user",role="user",content=message)
    x = agent(msg)
    history.append((message, x.content))
    return x.content

# 创建 Gradio 界面
iface = gr.ChatInterface(
    fn=chat,
    title="AI 助手",
    description="与 AI 助手进行对话。输入 'exit' 结束对话。",
    examples=["你好", "帮我查询系统版本", "写一个定时1小时后关机的 python 脚本"],
)

# 启动 Gradio 应用
iface.launch()